import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
from functools import lru_cache
from datetime import datetime
import re
import matplotlib.font_manager
import warnings
import sys

# Windows 控制台 UTF-8 设置（避免 emoji 编码错误）
if sys.platform == 'win32':
    try:
        import io
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace', line_buffering=True)
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace', line_buffering=True)
    except Exception:
        pass

# 抑制字体警告
warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib')

# 查找并设置中文字体
def setup_chinese_font():
    """自动检测并设置中文字体"""
    # Windows 常用中文字体列表（按优先级）
    chinese_fonts = [
        'Microsoft YaHei',      # 微软雅黑
        'SimHei',               # 黑体
        'SimSun',               # 宋体
        'KaiTi',                # 楷体
        'FangSong',             # 仿宋
        'STSong',               # 华文宋体
        'STKaiti',              # 华文楷体
        'Arial Unicode MS',     # macOS
    ]
    
    # 获取系统所有可用字体
    available_fonts = set([f.name for f in matplotlib.font_manager.fontManager.ttflist])
    
    # 查找第一个可用的中文字体
    for font in chinese_fonts:
        if font in available_fonts:
            print(f"使用字体: {font}")
            plt.rcParams['font.sans-serif'] = [font]
            plt.rcParams['axes.unicode_minus'] = False
            return font
    
    # 如果没有找到，使用默认设置并警告
    print("警告: 未找到合适的中文字体，图表中的中文可能无法正常显示")
    print(f"可用字体: {sorted(list(available_fonts)[:10])}...")
    return None

# 设置字体
setup_chinese_font()

DATA_PATH = Path("D:\\Quant\\quantdata\\stock-trading-data-pro")
BENCHMARK_PATH = Path("D:\\Quant\\quantdata\\stock-main-index-data")
CONCEPT_MAP_DIR = Path("D:\\Quant\\quantdata\\tdxexport\\alldata")  # 改为目录路径
OUTPUT_DIR = Path("./rrg_analysis_results")

RRG_N_PERIOD = 10
PLOT_TAIL_LENGTH = 4
BENCHMARK_CODE = 'sh000300'

# --- 核心功能函数 ---

def get_concept_file_by_date(target_date=None):
    """
    根据日期获取对应的tdx_blocks pkl文件
    
    Args:
        target_date: 目标日期字符串，格式：'YYYY-MM-DD'，如果为None则使用最新日期
    
    Returns:
        (file_path, date_str): 文件路径和日期字符串的元组
    """
    try:
        if target_date is None:
            # 使用最新日期
            return get_latest_concept_file()
        
        # 查找指定日期的文件
        target_file = CONCEPT_MAP_DIR / f'tdx_blocks_{target_date}.pkl'
        if target_file.exists():
            print(f"找到指定日期的概念映射文件: {target_file.name}")
            return target_file, target_date
        else:
            print(f"错误: 未找到日期为 {target_date} 的pkl文件: {target_file}")
            print(f"请检查文件是否存在，或运行 generate_blocks.py 生成该日期的数据")
            return None, None
            
    except Exception as e:
        print(f"查找指定日期概念文件时出错: {e}")
        return None, None

def get_latest_concept_file():
    """获取最新日期的tdx_blocks pkl文件，返回(file_path, date_str)"""
    try:
        # 查找所有 tdx_blocks_*.pkl 文件
        pattern = re.compile(r'tdx_blocks_(\d{4}-\d{2}-\d{2})\.pkl$')
        pkl_files = []
        
        for file_path in CONCEPT_MAP_DIR.glob('tdx_blocks_*.pkl'):
            match = pattern.search(file_path.name)
            if match:
                date_str = match.group(1)
                pkl_files.append((date_str, file_path))
        
        # 如果没有找到带日期的文件，尝试使用默认文件名
        if not pkl_files:
            default_file = CONCEPT_MAP_DIR / 'tdx_blocks.pkl'
            if default_file.exists():
                print(f"使用默认文件: {default_file}")
                return default_file, None  # 返回文件路径和日期
            else:
                print(f"错误: 未找到任何tdx_blocks pkl文件")
                return None, None
        
        # 按日期排序，获取最新的文件
        pkl_files.sort(key=lambda x: x[0], reverse=True)
        latest_date, latest_file = pkl_files[0]
        
        print(f"找到最新的概念映射文件: {latest_file.name} (日期: {latest_date})")
        return latest_file, latest_date  # 返回文件路径和日期
        
    except Exception as e:
        print(f"查找最新概念文件时出错: {e}")
        return None, None

def get_concept_stock_mapping(target_date=None):
    """
    获取概念与股票的映射关系
    
    Args:
        target_date: 目标日期字符串，格式：'YYYY-MM-DD'，如果为None则使用最新日期
    
    Returns:
        (concept_map, concept_date): 映射字典和日期字符串的元组
    """
    try:
        # 根据是否指定日期获取文件
        concept_file, concept_date = get_concept_file_by_date(target_date)
        if concept_file is None:
            return None, None
        
        df_gn = pd.read_pickle(concept_file)
        concept_map = df_gn.groupby('板块名称')['股票代码'].apply(list).to_dict()
        return concept_map, concept_date  # 返回映射和日期
    except Exception as e:
        print(f"加载概念映射时出错: {e}")
        return None, None

@lru_cache(maxsize=None)
def load_stock_data(stock_code, end_date=None):
    """
    加载单个股票的数据，并计算复权价
    
    Args:
        stock_code: 股票代码
        end_date: 结束日期，如果指定，只返回该日期及之前的数据
    """
    file_path = DATA_PATH / f"{stock_code}.csv"
    if not file_path.exists():
        return None
    
    try:
        df = pd.read_csv(file_path, encoding='gbk', skiprows=1)
        df.rename(columns={'交易日期': 'date', '收盘价': 'close', '前收盘价': 'prev_close'}, inplace=True)
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)
        
        # 如果指定了结束日期，过滤数据（截断到前一天）
        if end_date:
            end_date_obj = pd.to_datetime(end_date)
            df = df[df.index < end_date_obj]
            if df.empty:
                print(f"⚠️ 股票 {stock_code} 在 {end_date} 及之前无数据")
                return None
        
        df['adj_factor'] = (df['close'] / df['prev_close']).cumprod()
        df['adj_close'] = df['adj_factor']
        
        return df[['adj_close']].tail(100)
    except Exception as e:
        print(f"处理股票 {stock_code} 数据时出错: {e}")
        return None

def calculate_concept_timeseries(concept_name, concept_stocks, all_stock_data):
    """计算概念板块的价格时间序列数据"""
    concept_dfs = [all_stock_data[stock] for stock in concept_stocks if stock in all_stock_data and all_stock_data[stock] is not None]
    
    if not concept_dfs:
        return None

    combined_df = pd.concat(concept_dfs)
    daily_grouped = combined_df.groupby(combined_df.index)
    concept_price = daily_grouped['adj_close'].mean()
    
    result_df = pd.DataFrame({'price': concept_price}).dropna()
    return result_df

def load_benchmark_data(benchmark_code, end_date=None):
    """
    加载基准指数数据
    
    Args:
        benchmark_code: 基准指数代码
        end_date: 结束日期，如果指定，只返回该日期及之前的数据
    """
    file_path = BENCHMARK_PATH / f"{benchmark_code}.csv"
    if not file_path.exists():
        print(f"基准文件不存在: {file_path}")
        return None
    
    df = pd.read_csv(file_path)
    df.rename(columns={'candle_end_time': 'date'}, inplace=True)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    
    # 如果指定了结束日期，过滤数据（截断到前一天）
    if end_date:
        end_date_obj = pd.to_datetime(end_date)
        df = df[df.index < end_date_obj]
        if df.empty:
            print(f"⚠️ 基准指数 {benchmark_code} 在 {end_date} 及之前无数据")
            return None
    
    return df['close'] / df['close'].iloc[0]

def calculate_rrg(series, benchmark_series, n=14):
    """计算RRG指标"""
    aligned_series, aligned_benchmark = series.align(benchmark_series, join='inner')
    if aligned_series.empty:
        return None

    rs = aligned_series.rolling(window=int(n/2)).mean() / aligned_benchmark.rolling(window=int(n/2)).mean()
    rs = rs.dropna()

    rs_ratio = (rs - rs.rolling(window=n).mean()) / rs.rolling(window=n).std()
    rs_momentum = (rs_ratio - rs_ratio.rolling(window=n).mean()) / rs_ratio.rolling(window=n).std()

    rs_ratio = rs_ratio * 15 + 100
    rs_momentum = rs_momentum * 15 + 100

    data = pd.DataFrame({
        'RS-Ratio': rs_ratio,
        'RS-Momentum': rs_momentum
    }).dropna()
    
    return data

def plot_rrg(rrg_data, title, output_filename):
    """绘制RRG图并将相位数据保存到CSV"""
    fig, ax = plt.subplots(figsize=(16, 12))
    
    ax.axhline(100, color='grey', linestyle='--')
    ax.axvline(100, color='grey', linestyle='--')
    
    ax.text(0.98, 0.98, '领先 (Leading)', fontsize=12, color='green', ha='right', va='top', transform=ax.transAxes)
    ax.text(0.02, 0.98, '改进 (Improving)', fontsize=12, color='blue', ha='left', va='top', transform=ax.transAxes)
    ax.text(0.02, 0.02, '落后 (Lagging)', fontsize=12, color='red', ha='left', va='bottom', transform=ax.transAxes)
    ax.text(0.98, 0.02, '转弱 (Weakening)', fontsize=12, color='orange', ha='right', va='bottom', transform=ax.transAxes)

    phase_data_list = []

    for name, data in rrg_data.items():
        if data.empty:
            continue
        
        latest_data = data.iloc[-1]
        x, y = latest_data['RS-Ratio'], latest_data['RS-Momentum']

        # if not (x < 100 and y < 100): #剔除第三象限数据
        plot_data = data.tail(PLOT_TAIL_LENGTH)

        if len(plot_data) > 1:
            x_coords = plot_data['RS-Ratio']
            y_coords = plot_data['RS-Momentum']
            dx = x_coords.diff()
            dy = y_coords.diff()

            phases_rad = np.arctan2(dy, dx)
            phases_deg = np.degrees(phases_rad)
            for i in range(len(plot_data)):
                point_data = {
                    'concept': name,
                    'date': plot_data.index[i].strftime('%Y-%m-%d'),
                    'x_ratio': x_coords.iloc[i],
                    'y_momentum': y_coords.iloc[i],
                    'phase_rad': phases_rad.iloc[i],
                    'phase_deg': phases_deg.iloc[i]
                }
                phase_data_list.append(point_data)

            latest_phase_rad = phases_rad.iloc[-1]
            if not np.isnan(latest_phase_rad) and -np.pi / 2 <= latest_phase_rad <= np.pi / 2:
                ax.plot(plot_data['RS-Ratio'], plot_data['RS-Momentum'], marker='o', linestyle='-', markersize=3, alpha=0.6)
                ax.scatter(x, y, s=120, alpha=0.8)
                ax.text(x, y, name, fontsize=9)

    ax.set_xlabel('JdK RS-Ratio (相对强度)')
    ax.set_ylabel('JdK RS-Momentum (相对动量)')
    ax.set_title(title, fontsize=16)
    ax.grid(True)
    plt.tight_layout()
    
    plt.savefig(output_filename)
    print(f"RRG图已保存至: {output_filename}")
    plt.close(fig)

    if phase_data_list:
        phase_df = pd.DataFrame(phase_data_list)
        csv_filename = output_filename.with_suffix('.csv')
        phase_df.to_csv(csv_filename, index=False, encoding='utf-8-sig')
        print(f"相位数据已保存至: {csv_filename}")

def main(target_date=None):
    """
    主执行函数
    
    Args:
        target_date: 目标日期字符串，格式：'YYYY-MM-DD'，如果为None则使用最新日期
    """
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    print("="*80)
    if target_date:
        print(f"指定日期模式: {target_date}")
    else:
        print("自动模式: 使用最新日期")
    
    print("正在加载概念映射...")
    concept_stock_map, concept_date = get_concept_stock_mapping(target_date)
    if not concept_stock_map:
        print("未能加载概念映射，程序退出。")
        return
    
    print(f"成功加载 {len(concept_stock_map)} 个概念板块。")
    
    # 显示数据日期信息
    if concept_date:
        print(f"📅 概念板块数据日期: {concept_date}")
        print(f"⚠️  股票和基准指数数据将被截断到 {concept_date} 的前一天（不包含 {concept_date}）")
        date_str = concept_date  # 使用概念数据的日期
    else:
        print("⚠️  未能获取概念数据日期，将使用所有可用数据")
        date_str = datetime.now().strftime('%Y-%m-%d')
    
    print("="*80)

    all_stocks_in_concepts = sorted(list(set(stock for stocks in concept_stock_map.values() for stock in stocks)))
    print(f"共涉及 {len(all_stocks_in_concepts)} 只股票，开始加载数据...")
    
    all_stock_data = {}
    for i, stock_code in enumerate(all_stocks_in_concepts):
        if (i + 1) % 200 == 0:
            print(f"  已加载 {i + 1}/{len(all_stocks_in_concepts)}...")
        # 传入concept_date作为结束日期
        data = load_stock_data(stock_code, end_date=concept_date)
        if data is not None and not data.empty:
            all_stock_data[stock_code] = data
    print("所有股票数据加载完成。")
    load_stock_data.cache_clear()

    print(f"正在加载基准指数 {BENCHMARK_CODE}...")
    # 传入concept_date作为结束日期
    benchmark_series = load_benchmark_data(BENCHMARK_CODE, end_date=concept_date)
    if benchmark_series is None:
        print(f"未能加载基准指数 {BENCHMARK_CODE}，程序退出。")
        return
    print("基准数据加载完成。")

    all_price_rrg = {}
    
    print("开始遍历概念并计算RRG指标...")
    for i, (concept_name, stocks) in enumerate(concept_stock_map.items()):
        print(f"  处理中 ({i+1}/{len(concept_stock_map)}): {concept_name}")
        
        concept_ts = calculate_concept_timeseries(concept_name, stocks, all_stock_data)
        
        if concept_ts is None or concept_ts.empty:
            print(f"    -> 跳过，无有效数据")
            continue
            
        price_rrg = calculate_rrg(concept_ts['price'], benchmark_series, n=RRG_N_PERIOD)
        if price_rrg is not None:
            all_price_rrg[concept_name] = price_rrg
            
    print("所有概念处理完毕。")

    print("打印各象限数据：")
    if all_price_rrg:
        for name, data in all_price_rrg.items():
            if data.empty:
                continue
            latest_data = data.iloc[-1]
            x, y = latest_data['RS-Ratio'], latest_data['RS-Momentum']
            if (x < 100 and y > 100):
                print(f"  ✓ 改进象限: {name} (RS-Ratio: {x:.1f}, RS-Momentum: {y:.1f})")
        for name, data in all_price_rrg.items():
            if data.empty:
                continue
            latest_data = data.iloc[-1]
            x, y = latest_data['RS-Ratio'], latest_data['RS-Momentum']
            if (x >  100 and y > 100):
                print(f"  ✓ 领先象限: {name} (RS-Ratio: {x:.1f}, RS-Momentum: {y:.1f})")
        for name, data in all_price_rrg.items():
            if data.empty:
                continue
            latest_data = data.iloc[-1]
            x, y = latest_data['RS-Ratio'], latest_data['RS-Momentum']
            if (x >  100 and y < 100):
                print(f"  ✓ 转弱象限: {name} (RS-Ratio: {x:.1f}, RS-Momentum: {y:.1f})")

    print("正在生成RRG图...")
    if all_price_rrg:
        price_title = f'概念板块价格相对轮动图 (基准: {BENCHMARK_CODE}) - {date_str}'
        price_filename = OUTPUT_DIR / f'concept_price_RRG_{date_str}.png'
        plot_rrg(all_price_rrg,
                 title=price_title,
                 output_filename=price_filename)
    else:
        print("没有足够的数据生成价格RRG图。")
        
    print("分析完成。")


if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(
        description='RRG相对轮动图分析工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
使用示例:
  # 使用最新日期的数据（默认）
  python rgg_ana_lite.py
  
  # 指定日期生成分析
  python rgg_ana_lite.py --date 2025-12-04
  python rgg_ana_lite.py -d 2025-12-04
  
注意:
  - 日期格式必须为 YYYY-MM-DD
  - 需要先运行 generate_blocks.py 生成对应日期的 pkl 文件
  - pkl 文件路径: D:\\Quant\\quantdata\\tdxexport\\alldata\\tdx_blocks_YYYY-MM-DD.pkl
        ''')
    
    parser.add_argument(
        '-d', '--date',
        type=str,
        default=None,
        metavar='YYYY-MM-DD',
        help='指定要分析的日期（格式：YYYY-MM-DD），不指定则使用最新日期'
    )
    
    args = parser.parse_args()
    
    # 验证日期格式
    if args.date:
        import re
        if not re.match(r'^\d{4}-\d{2}-\d{2}$', args.date):
            print(f"错误: 日期格式不正确: {args.date}")
            print("正确格式: YYYY-MM-DD，例如: 2025-12-04")
            exit(1)
    
    main(target_date=args.date)